MODULE fem_A_matrix
      
   USE fem_global_types
   USE fem_global_consts
   USE fem_A_matrix_LIP

   IMPLICIT NONE
   PRIVATE
   ! Public procedures
   PUBLIC :: fem_init_A_matrix,      &
             fem_free_A_matrix

   ! Public variables
   PUBLIC :: Amat_initialized,       &
             Amat,                   &
             Amat_dia_inv

   REAL(REALK), ALLOCATABLE, TARGET, SAVE :: Amat(:,:,:,:,:,:)
   REAL(REALK), ALLOCATABLE, TARGET, SAVE :: Amat_dia_inv(:)
   LOGICAL, SAVE :: Amat_initialized = .FALSE.

CONTAINS

!------------------------------------------------------------------------------

   SUBROUTINE fem_init_A_matrix(scheme)
   
      IMPLICIT NONE
      TYPE(fem_scheme), INTENT(IN) :: scheme

      REAL(REALK), ALLOCATABLE :: A0mat(:,:), A1mat(:,:) 
      INTEGER(INTK) :: I, J, K, II, JJ, KK, n

      n = scheme%FEFNNod1D

      ALLOCATE(Amat(n,n,n,n,n,n))
      ALLOCATE(A1mat(n,n))
      ALLOCATE(A0mat(n,n))
      
      SELECT CASE (scheme%FEFNthShp)
      CASE(1)
         A1mat(:,:) = RESHAPE(A1mat_1st, (/n, n/))
         A0mat(:,:) = RESHAPE(A0mat_1st, (/n, n/))
      CASE(2)
         A1mat(:,:) = RESHAPE(A1mat_2nd, (/n, n/))
         A0mat(:,:) = RESHAPE(A0mat_2nd, (/n, n/))
      CASE(3)
         A1mat(:,:) = RESHAPE(A1mat_3rd, (/n, n/))
         A0mat(:,:) = RESHAPE(A0mat_3rd, (/n, n/))
      CASE(4)
         A1mat(:,:) = RESHAPE(A1mat_4th, (/n, n/))
         A0mat(:,:) = RESHAPE(A0mat_4th, (/n, n/))
      CASE(5)
         A1mat(:,:) = RESHAPE(A1mat_5th, (/n, n/))
         A0mat(:,:) = RESHAPE(A0mat_5th, (/n, n/))
      END SELECT

!FIXME allow for non-cubic 3D elements!
      A1mat(:,:) = A1mat(:,:) / scheme%FEFNodEdg(1)
      A0mat(:,:) = A0mat(:,:) * scheme%FEFNodEdg(1)

      DO KK = 1, n
      DO JJ = 1, n
      DO II = 1, n
         DO K = 1, n
         DO J = 1, n
         DO I = 1, n
            Amat(I,J,K,II,JJ,KK) = A1mat(I,II) * A0mat(J,JJ) * A0mat(K,KK)  &
                                 + A0mat(I,II) * A1mat(J,JJ) * A0mat(K,KK)  &
                                 + A0mat(I,II) * A0mat(J,JJ) * A1mat(K,KK) 
         END DO
         END DO
         END DO
      END DO
      END DO
      END DO

      DEALLOCATE(A1mat, A0mat)
      Amat_initialized = .TRUE.

   END SUBROUTINE fem_init_A_matrix

!-------------------------------------------------------------------------------

   SUBROUTINE fem_free_A_matrix
     
      IMPLICIT NONE
      IF (ALLOCATED(Amat)) DEALLOCATE(Amat)
      IF (ALLOCATED(Amat_dia_inv)) DEALLOCATE(Amat_dia_inv)
      Amat_initialized = .FALSE.

   END SUBROUTINE fem_free_A_matrix

!------------------------------------------------------------------------------

END MODULE fem_A_matrix
